import pandas as pd

def compute_retention(df: pd.DataFrame, days: int = 7) -> pd.DataFrame:
    """
    计算 1 到 N 天的用户留存率。
    """
    df['date'] = pd.to_datetime(df['timestamp']).dt.normalize()  # 确保是datetime64[ns]

    user_first_visit = df.groupby('visitorid')['date'].min().reset_index()
    user_first_visit.columns = ['visitorid', 'reg_date']

    df = df.merge(user_first_visit, on='visitorid')
    df['day_diff'] = (df['date'] - df['reg_date']).dt.days
    df = df[df['day_diff'] <= days]

    retention = df.groupby(['reg_date', 'day_diff'])['visitorid'].nunique().unstack(fill_value=0)
    retention_rate = retention.divide(retention[0], axis=0).round(4)
    retention_rate['total_users'] = retention[0]

    print(f"📊 留存计算完成，共 {retention_rate.shape[0]} 个注册日")
    return retention_rate.reset_index()

def generate_labels(df: pd.DataFrame, retention_days: int = 3) -> pd.DataFrame:
    df['event_time'] = pd.to_datetime(df['timestamp'])
    df['datetime'] = df['event_time']

    user_first_date = df.groupby('visitorid')['datetime'].min()
    user_last_date = df.groupby('visitorid')['datetime'].max()

    label_df = pd.DataFrame({
        'visitorid': user_first_date.index,
        'label': ((user_last_date - user_first_date).dt.days >= retention_days).astype(int)
    })

    print(f"✅ 留存标签生成完成，共 {label_df.shape[0]} 个用户")
    return label_df
